from db.db_customer import DBCustomer
from db.db_bill import DBBill
from db.db_sale import DBSale


__all__ = ["db_cust_wrapper", "db_cust_and_bill_wrapper", "db_cust_and_sale_wrapper"]


def db_cust_wrapper(method):
    def wrapper(self, *args, **kwargs):
        with DBCustomer() as _cust_db:
            result = method(self, cust_db=_cust_db,  **kwargs)
        return result
    return wrapper

def db_cust_and_bill_wrapper(method):
    def wrapper(self, *args, **kwargs):
        with DBCustomer() as _cust_db:
            with DBBill() as bill_db:
                result = method(self, cust_db=_cust_db, bill_db=bill_db,  **kwargs)
        return result
    return wrapper

def db_cust_and_sale_wrapper(method):
    def wrapper(self, *args, **kwargs):
        with DBCustomer() as _cust_db:
            with DBSale() as _sale_db:
                result = method(self, cust_db=_cust_db, sale_db=_sale_db,  **kwargs)
        return result
    return wrapper